package com.ollama.study.controller;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.converter.ListOutputConverter;
import org.springframework.ai.converter.MapOutputConverter;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.convert.support.DefaultConversionService;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.beans.factory.annotation.Value;

import javax.print.attribute.standard.Media;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.*;

import java.util.List;
import java.util.Map;
import java.util.logging.Logger;

/*
 * ollama的结构化输出设置
 * 此处脚本的主要目的是实现对调用大模型的结构化输出的相关调整操作
 * */

@RestController
@RequestMapping("/ollamaStructuredOutput")
@Tag(name = "ollama结构化输出")
public class OllamaStructuredOutputController {

    // 创建一个日志记录器,主要的作用是获取ollama的日志
    private static final Logger log = Logger.getLogger(OllamaStructuredOutputController.class.getName());

    @Resource(name = "ollamaChatModelD1")
    private OllamaChatModel ollamaChatModel;

    @Value("${spring.ai.ollama.base-url}")
    private String OLLAMA_BASE_URL;// 获取配置文件中的ollama的base-url

    @Value("${spring.ai.ollama.chat.model}")
    private String OLLAMA_MODEL_NAME;// 获取配置文件中的ollama的chat.model

    @GetMapping("/list")
    public List<String> listOutputConverter(@RequestParam String message) {
        DefaultConversionService conversionService = new DefaultConversionService();// 接收请求参数message，类型是String
        ListOutputConverter outputConverter = new ListOutputConverter(conversionService);// 创建一个Spring的转换服务实例，用于类型转换
        // 创建一个自定义的ListOutputConverter实例，传入转换服务 该类用于将模型输出的文本转换为List<String>
        String format = outputConverter.getFormat();
        // 获取ListOutputConverter定义的格式字符串，用于提示模板
        String template = """
        推荐五个 {subject}
        {format}
        """;
        // 定义一个多行字符串模板，用于生成提示词
        // {subject} 和 {format} 是占位符，会被替换成具体内容
        PromptTemplate promptTemplate = PromptTemplate.builder()
                .template(template)
                .variables(Map.of("subject", message, "format", format))
                .build();
        // 使用Builder模式创建一个PromptTemplate对象
        // template是模板字符串
        // variables用来替换模板中的占位符，subject替换为请求参数message，format替换为格式字符串
        ChatResponse response = ollamaChatModel.call(
                new Prompt(promptTemplate.createMessage(),
                        OllamaOptions.builder()
                                .model(OLLAMA_MODEL_NAME)
                                .temperature(0.4)
                                .build()
                ));
        // 调用chatModel的call方法发送Prompt，得到模型返回的响应
        // Prompt中包含了生成的消息
        // OllamaOptions设置了模型名和温度参数（温度控制生成的随机性）
        log.info("mapOutputConverter - 模型响应内容: " + response.getResult().getOutput().getText());
        // 记录日志，打印模型返回的文本结果
        List<String> list = outputConverter.convert(response.getResult().getOutput().getText());
        // 使用outputConverter将模型返回的文本转换为List<String>
        return list;
        // 返回转换后的字符串列表
    }

    @GetMapping("/map")
    public Map<String, Object> mapOutputConverter(@RequestParam String message) {
        MapOutputConverter outputConverter = new MapOutputConverter();
        String format = outputConverter.getFormat();
        String template = """
				请推荐一首诗人{author}写的最受欢迎的诗，要包含作者名称author，作者所属朝代year，诗的标题theme，诗的内容context：
				{format}
				""";
        PromptTemplate promptTemplate = PromptTemplate.builder()
                .template(template)
                .variables(Map.of("author", message, "format", format))
                .build();
        ChatResponse response = ollamaChatModel.call(
                new Prompt(promptTemplate.createMessage(),
                        OllamaOptions.builder()
                                .model(OLLAMA_MODEL_NAME)
                                .temperature(0.4)
                                .build()
                ));
        log.info("mapOutputConverter - 模型响应内容: " + response.getResult().getOutput().getText());
        Map<String, Object> result = outputConverter.convert(response.getResult().getOutput().getText());
        return result;
    }
}
